/**
 * Copyright (c) 2024 Hunan OpenValley Digital Industry Development Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "GPUImageWhiteBalanceFilter.h"

const std::string GPUImageWhiteBalanceFilter::WHITE_BALANCE_FRAGMENT_SHADER =
    R"(
    #extension GL_OES_EGL_image_external : require
    uniform samplerExternalOES inputImageTexture;
    varying highp vec2 textureCoordinate;

    uniform lowp float temperature;
    uniform lowp float tint;

    const lowp vec3 warmFilter = vec3(0.93, 0.54, 0.0);

    const mediump mat3 RGBtoYIQ = mat3(
        0.299,     0.587,     0.114,
        0.596,    -0.274,    -0.322,
        0.212,    -0.523,     0.311
    );

    const mediump mat3 YIQtoRGB = mat3(
        1.0,  0.956,     0.621,
        1.0, -0.272,    -0.647,
        1.0, -1.105,     1.702
    );

    void main()
    {
        lowp vec4 source = texture2D(inputImageTexture, textureCoordinate);

        mediump vec3 yiq = RGBtoYIQ * source.rgb; // adjusting tint
        yiq.b = clamp(yiq.b + tint * 0.5226 * 0.1, -0.5226, 0.5226);
        lowp vec3 rgb = YIQtoRGB * yiq;

        lowp vec3 processed = vec3(
            (rgb.r < 0.5 ? (2.0 * rgb.r * warmFilter.r) : (1.0 - 2.0 * (1.0 - rgb.r) * (1.0 - warmFilter.r))), // adjusting temperature
            (rgb.g < 0.5 ? (2.0 * rgb.g * warmFilter.g) : (1.0 - 2.0 * (1.0 - rgb.g) * (1.0 - warmFilter.g))),
            (rgb.b < 0.5 ? (2.0 * rgb.b * warmFilter.b) : (1.0 - 2.0 * (1.0 - rgb.b) * (1.0 - warmFilter.b)))
        );

        gl_FragColor = vec4(mix(rgb, processed, temperature), source.a);
    }
)";

GPUImageWhiteBalanceFilter::GPUImageWhiteBalanceFilter()
    : temperature(5000.0f), tint(0.0f) {}

GPUImageWhiteBalanceFilter::GPUImageWhiteBalanceFilter(
    const float temperatureVal, const float tintVal)
    : GPUImageFilter(NO_FILTER_VERTEX_SHADER, WHITE_BALANCE_FRAGMENT_SHADER)
{
    temperature = temperatureVal;
    tint = tintVal;
}

void GPUImageWhiteBalanceFilter::onInit()
{
    GPUImageFilter::onInit();
    temperatureLocation = glGetUniformLocation(getProgram(), "temperature");
    tintLocation = glGetUniformLocation(getProgram(), "tint");
}

void GPUImageWhiteBalanceFilter::onInitialized()
{
    GPUImageFilter::onInitialized();
    setTemperature(temperature);
    setTint(tint);
}

void GPUImageWhiteBalanceFilter::setTemperature(const float temperatureVal)
{
    temperature = temperatureVal;
    setFloat(temperatureLocation, (temperature < 5000.0f)
                                      ? (0.0004f * (temperature - 5000.0f))
                                      : (0.00006f * (temperature - 5000.0f)));
}

void GPUImageWhiteBalanceFilter::setTint(const float tintVal)
{
    tint = tintVal;
    setFloat(tintLocation, tint / 100.0f);
}